import bst  # see bst.py from April 22-28


red = 'red'
black = 'black'


class RedBlackTree(bst.Tree):

    def __init__(self):
        self.root = None

    def insert(self, k, parent=None):
        node = super(RedBlackTree, self).insert(k)
        node.color = red
        x = node
        while x != self.root and x.p.color == red:
            if x.p == x.p.p.left:
                if RedBlackTree._recolor(x.p, x.p.p.right):
                    x = x.p.p
                else:
                    if x == x.p.right:
                        x = x.p
                        self._left_rotate(x)
                    x.p.color = black
                    x.p.p.color = red
                    self._right_rotate(x.p.p)
            else:
                y = x.p.p.left
                if RedBlackTree._recolor(x.p, x.p.p.left):
                    x = x.p.p
                else:
                    if x == x.p.left:
                        x = x.p
                        self._right_rotate(x)
                    x.p.color = black
                    x.p.p.color = red
                    self._left_rotate(x.p.p)

        self.root.color = black
        return node

    @staticmethod
    def _recolor(father, uncle):
        if uncle and uncle.color == red:
            father.color = uncle.color = black
            father.p.color = red
            return True
        return False

    def _left_rotate(self, x):
        y = x.right

        x.right = y.left
        if x.right:
            x.right.p = x

        y.p = x.p
        if y.p:
            if x == x.p.left:
                x.p.left = y
            else:
                x.p.right = y
        else:
            self.root = y

        y.left = x
        x.p = y

    def _right_rotate(self, x):
        y = x.left

        x.left = y.right
        if x.left:
            x.left.p = x

        y.p = x.p
        if y.p:
            if x == x.p.left:
                x.p.left = y
            else:
                x.p.right = y
        else:
            self.root = y

        y.right = x
        x.p = y


